package com.ht.module.ai.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import com.ht.module.ai.config.WeatherTools;
import com.ht.module.ai.config.Assistant;
import com.ht.module.ai.config.OllamaConfig;
import com.ht.module.ai.config.PersistentChatMemoryStore;
import com.ht.module.ai.entity.AiChatSession;
import com.ht.module.ai.service.AiChatSessionService;
import com.ht.module.ai.service.AiService;
import com.ht.module.ai.vo.AiChatVo;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.service.AiServices;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import reactor.core.publisher.Flux;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;


@Slf4j
@Service
public class AiServiceImpl implements AiService {
    @Resource
    private OllamaConfig ollamaConfig;
    @Resource
    private AiChatSessionService aiChatSessionService;
    @Resource
    private PersistentChatMemoryStore persistentChatMemoryStore;
    @Resource
    private WeatherTools weatherTools;

    @Override
    @Transactional
    public Flux<String> streamMemoryChat(AiChatVo vo) {
        //查询处当前会话
        AiChatSession session = aiChatSessionService.getById(vo.getSessionId());

        //判断是否存在会话id
        if (BeanUtil.isEmpty(session)) {
            //不存在，创建新的会话
            session = aiChatSessionService.createSession(vo.getSessionId(), vo.getMsg());
        }

        OllamaStreamingChatModel model = OllamaStreamingChatModel.builder()
                .baseUrl(ollamaConfig.baseUrl)
                .modelName(ollamaConfig.modelName)
                .build();
        ChatMemory chatMemory = null;
        if (StrUtil.isNotEmpty(vo.getSessionId())) {
            chatMemory = MessageWindowChatMemory.builder()
                    .id(vo.getSessionId())
                    .maxMessages(10)
                    .chatMemoryStore(persistentChatMemoryStore)
                    .build();
        }
        //构建会话service
        AiServices<Assistant> builder = AiServices.builder(Assistant.class)
                //引入工具后会阻断流式输出
                .tools(weatherTools)
                .streamingChatLanguageModel(model);
        //存在session，才进行记忆会话
        if (StrUtil.isNotEmpty(vo.getSessionId())) {
            builder.chatMemory(chatMemory);
        }

        Assistant assistant = builder.build();
        //定义ai回复的消息内容集合
        List<String> msgs = new ArrayList<>();
        // 在流结束时执行某个操作

        return assistant.chat(vo.getMsg())
                //获取回复内容片段
                .doOnNext(msgs::add)
                .doOnComplete(() -> {
                    String fullMessgae = String.join("", msgs);
                    log.info("ai最终回复内容：{}", fullMessgae);
                });
    }


}
